#include "txr_algo_dlm_v1.h"
#include "darknet/include/tk_dlm.hpp"
#include <vector>
namespace txr_algo_dlm
{

TxrAlgoDlmV1::TxrAlgoDlmV1()
	:m_p_dlm(TK_NULL_PTR)
{

}

TxrAlgoDlmV1::~TxrAlgoDlmV1()
{
	Release();
}

tk_bool TxrAlgoDlmV1::Load(st_dlm_cfg & cfg)
{
	Release();
	std::string path_cfg = std::string(cfg.file_path_name) + ".dat";
	std::string path_weight = std::string(cfg.file_path_name) + ".dats";

	m_p_dlm = new tk_dlm();
	return m_p_dlm->Load(path_cfg.c_str(), path_weight.c_str(), cfg.gpu_id);
}

tk_int32 TxrAlgoDlmV1::BatchSize()
{
	return 1;
}

void TxrAlgoDlmV1::Release()
{
	SAFE_DELETE(m_p_dlm);
}

void BoxTrans(bbox_t * p_box, tk_int32 box_num, st_result_boxes & result)
{
	result.vaild_num = box_num;
	for (tk_int32 i = 0; i < box_num; ++i)
	{
		result.boxes[i].x = p_box[i].x;
		result.boxes[i].y = p_box[i].y;
		result.boxes[i].w = p_box[i].w;
		result.boxes[i].h = p_box[i].h;
		result.boxes[i].prob = p_box[i].prob;
		result.boxes[i].obj_id = p_box[i].obj_id;
	}
}

void TxrAlgoDlmV1::Detect(st_detect_unit * p_unit, tk_int32 size)
{
	std::vector<bbox_t> v_box;
	v_box.resize(C_SHARP_MAX_OBJECTS);
	tk_int32 box_num = 0;
	for (tk_int32 i = 0; i < size; ++i)
	{
		box_num = 0;
		st_img_rgb & rgb = p_unit[i].img;
		if (rgb.pp_rgbf[0] && rgb.pp_rgbf[1] && rgb.pp_rgbf[2])
		{
			image_rgbf_t img_t;
			img_t.data = rgb.pp_rgbf;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}
		else if(rgb.pp_rgb8[0] && rgb.pp_rgb8[1] && rgb.pp_rgb8[2])
		{
			image_rgb8_t img_t;
			img_t.data = rgb.pp_rgb8;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}
		else if (rgb.pp_rgb16[0] && rgb.pp_rgb16[1] && rgb.pp_rgb16[2])
		{
			image_rgb16_t img_t;
			img_t.data = rgb.pp_rgb16;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}
		else if (rgb.pp_dev_rgbf[0] && rgb.pp_dev_rgbf[1] && rgb.pp_dev_rgbf[2])
		{
			image_rgbf_t img_t;
			img_t.data_dev = rgb.pp_dev_rgbf;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}
		else if (rgb.pp_dev_rgb16[0] && rgb.pp_dev_rgb16[1] && rgb.pp_dev_rgb16[2])
		{
			image_rgb16_t img_t;
			img_t.data_dev = rgb.pp_dev_rgb16;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}
		else if (rgb.pp_dev_rgb8[0] && rgb.pp_dev_rgb8[1] && rgb.pp_dev_rgb8[2])
		{
			image_rgb8_t img_t;
			img_t.data_dev = rgb.pp_dev_rgb8;
			img_t.h = rgb.h;
			img_t.w = rgb.w;
			img_t.c = rgb.c;
			box_num = m_p_dlm->Detect(img_t, v_box.data(), 0.01);
		}

		BoxTrans(v_box.data(), box_num,p_unit[i].result);
	}
}
//=======================================================
extern "C" TK_TXR_ALGO_DLM_V1_EXPORT TxrAlgoDlmImpl* CreateDlmObj();
extern "C" TK_TXR_ALGO_DLM_V1_EXPORT void DeleteDlmObj(TxrAlgoDlmImpl* p);

TxrAlgoDlmImpl* CreateDlmObj()
{
	return new TxrAlgoDlmV1();
}

void DeleteDlmObj(TxrAlgoDlmImpl* p)
{
	SAFE_DELETE(p);
}

}//namespace txr_algo_dlm